'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-04-01 15:21:54
'''
# -*- coding: UTF-8 -*-
from subprocess import *
import random, time, json, sys
# from kafka import KafkaProducer
# from kafka import KafkaConsumer
# from kafka import TopicPartition
# import sys
# sys.path.append('../util')
# from feature_handler import get_fea_index, fea_code


class BatchTrain():
    def __init__(self, model_name, batch_size, shuffle_size, watch_num, neg_rate, process_num, ps_url, dates, f_path, run_shell, net_conf, labels, random_sleep=0):
        # model_name += ('_' + str(time.strftime("%Y%m%d_%H%M%S", time.localtime())))
        self.model_name = model_name
        print("model_name = ", model_name)
        # exit(0)
        self.batch_size = batch_size
        self.process_num = process_num
        self.shuffle_size = shuffle_size
        self.watch_num = watch_num
        self.neg_rate = neg_rate
        self.run_shell = run_shell
        self.dates = dates
        self.f_path = f_path
        self.random_sleep = random_sleep
        self.net_conf = net_conf
        self.labels = labels
        self.ps_url = ps_url
        self.shell_p = self.run_shell + " -nets /data1/xuejian/sync/model_zoo/models/cascade/" + self.net_conf + " -mn " + self.model_name + " -w " + str(self.watch_num) + " -s " + str(self.shuffle_size) + " -b " + str(self.batch_size) + " -ne " + str(self.neg_rate) + " -labels " + self.labels + " -ps_url " + self.ps_url + ':'+str(sys.argv[1]) #+ ' -is_json true'

        print("shell = ", self.shell_p)



    def train(self):
        for date in self.dates:
            print(date)
            proc_dict = {}
            num_list = [0] * self.process_num
            print("num_list = ", num_list)
            for i in range(self.process_num):
                # print(i)
                f_date = self.f_path + date + '/' + str(i)
                # f_date = self.f_path + date + '/' + str(sys.argv[2])
                if i != 0:
                    time.sleep(random.randint(0, self.random_sleep))
                proc = Popen('cat ' + f_date + ' | ' + self.shell_p, shell=True)
                # print('process start')
                proc_dict[i] = proc
            for index, p in proc_dict.items():
                p.wait()
                print('finished!! : ', index)

    def train2(self):

        proc = Popen(self.shell_p, shell=True, stdin=PIPE)
        print('process start')
        for date in self.dates:
            print(date)
            with open(self.f_path + date, 'rb') as f_date:
                print(self.f_path + date)
                line = f_date.readline()
                while line != b'':
                    if line == b'':
                        break
                    proc.stdin.write(line)
                    proc.stdin.flush()
                    line = f_date.readline()

        proc.stdin.close()
        proc.wait()


if __name__ == '__main__':
    dates = []
    # dates.extend(['2020-04-' + str(i) for i in range(28, 31)])
    # dates.extend(['2020-05-0' + str(i) for i in range(1, 10)])
    # dates.extend(['2020-05-' + str(i) for i in range(10, 32)])
    # dates.extend(['2020-06-0' + str(i) for i in range(1, 6)])
    dates.extend(['2020-06-0' + str(i) for i in range(6, 8)])
    ps_url = "172.21.39.6" #offlineps-06 port 1111 slave 1112 deepfm_cascade
    # cut_path = '/data3/training_data/arc_six/bossrec_success_fid_flow_cut/'
    cut_path = '/data3/training_data/galaxy_data_cut/'
    single_path = '/data3/training_data/galaxy_data_new/'
    nums = 6
    if nums <= 1:
        real_path = single_path
        watch_num = 200000
    else:
        real_path = cut_path
        watch_num = 50000
    batch_train = BatchTrain('deepfm_train', 1000, 10000, watch_num, 1.0, nums, ps_url, dates, real_path, '/data1/xuejian/sync/flash_train/flash_train', 'deepfm.net', 'one', 14)
    batch_train.train()